{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eb681984-e50d-4c5e-b827-98618fa0c55c",
   "metadata": {},
   "source": [
    "# Surya\n",
    "Surya is a document OCR toolkit that does:\n",
    "\n",
    "OCR in 90+ languages that benchmarks favorably vs cloud services\n",
    "Line-level text detection in any language\n",
    "Layout analysis (table, image, header, etc detection)\n",
    "Reading order detection\n",
    "\n",
    "https://github.com/VikParuchuri/surya.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "3337d0f1-cda1-45ca-a2f3-408e51d11ee1",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/laobao/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from surya.input.load import load_from_file, load_from_folder, load_lang_file\n",
    "from surya.input.langs import replace_lang_with_code, get_unique_langs\n",
    "\n",
    "from surya.model.detection.segformer import load_model as load_detection_model, load_processor as load_detection_processor\n",
    "from surya.model.recognition.model import load_model as load_recognition_model\n",
    "from surya.model.recognition.processor import load_processor as load_recognition_processor\n",
    "from surya.model.recognition.tokenizer import _tokenize\n",
    "from surya.ocr import run_ocr\n",
    "from surya.postprocessing.text import draw_text_on_image\n",
    "from surya.settings import settings\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "01520ea8-83c8-4e17-9a43-4ce33c281d13",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_path = './sample-docs/Doc2.pdf'\n",
    "max_pages = 10\n",
    "start_page = 0\n",
    "\n",
    "# 将PDF文件转换为image\n",
    "images, names = load_from_file(input_path, max_pages, start_page)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3545e773-c662-4ef7-8dc1-f296e849f0c1",
   "metadata": {},
   "source": [
    "## 检测模型和处理器"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4fc936ac-a710-4d85-b9ec-0ac3dd4b98b9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading detection model vikp/surya_det2 on device cpu with dtype torch.float32\n"
     ]
    }
   ],
   "source": [
    "det_processor = load_detection_processor()\n",
    "det_model = load_detection_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccfc4b06-1b63-4d21-91ff-4068a4503882",
   "metadata": {},
   "source": [
    "## 使用语言"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d1fece7d-9c95-423f-925d-ea4a88608839",
   "metadata": {},
   "outputs": [],
   "source": [
    "langs = \"zh\"\n",
    "langs = langs.split(\",\")\n",
    "replace_lang_with_code(langs)\n",
    "image_langs = [langs] * len(images)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8b14821a-057f-413c-9da6-57c773116175",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading recognition model vikp/surya_rec on device mps with dtype torch.float16\n"
     ]
    }
   ],
   "source": [
    "_, lang_tokens = _tokenize(\"\", get_unique_langs(image_langs))\n",
    "rec_model = load_recognition_model(langs=lang_tokens) # Prune model moe layer to only include languages we need\n",
    "rec_processor = load_recognition_processor()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "539eab50-0936-4264-ac1d-0f43c38c71ea",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Detecting bboxes: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:05<00:00,  5.52s/it]\n",
      "Recognizing Text:   0%|                                                                                                                                                                                    | 0/1 [00:00<?, ?it/s]/Users/laobao/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/transformers/generation/utils.py:1518: UserWarning: You have modified the pretrained model configuration to control generation. This is a deprecated strategy to control generation and will be removed soon, in a future version. Please use and modify the model generation configuration (see https://huggingface.co/docs/transformers/generation_strategies#default-text-generation-configuration )\n",
      "  warnings.warn(\n",
      "Recognizing Text: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:30<00:00, 30.67s/it]\n"
     ]
    }
   ],
   "source": [
    "predictions_by_image = run_ocr(images, image_langs, det_model, det_processor, rec_model, rec_processor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "46717ef1-adc2-41f1-b2ef-92958ed94529",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from collections import defaultdict\n",
    "\n",
    "result_path = './results'\n",
    "\n",
    "for idx, (name, image, pred, langs) in enumerate(zip(names, images, predictions_by_image, image_langs)):\n",
    "            bboxes = [l.bbox for l in pred.text_lines]\n",
    "            pred_text = [l.text for l in pred.text_lines]\n",
    "            page_image = draw_text_on_image(bboxes, pred_text, image.size, langs, has_math=\"_math\" in langs)\n",
    "            page_image.save(os.path.join(result_path, f\"{name}_{idx}_text.png\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "14bed070-9466-40d7-945d-c5db2e9dfb14",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_preds = defaultdict(list)\n",
    "for name, pred, image in zip(names, predictions_by_image, images):\n",
    "    out_pred = pred.model_dump()\n",
    "    out_pred[\"page\"] = len(out_preds[name]) + 1\n",
    "    out_preds[name].append(out_pred)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44c39220-a517-4874-88ec-f48b605d1008",
   "metadata": {},
   "source": [
    "## 识别输出为json文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6cf4618b-9b19-4786-ab45-d3a1cf7af595",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wrote results to ./results\n"
     ]
    }
   ],
   "source": [
    "with open(os.path.join(result_path, \"results.json\"), \"w+\", encoding=\"utf-8\") as f:\n",
    "    json.dump(out_preds, f, ensure_ascii=False)\n",
    "\n",
    "print(f\"Wrote results to {result_path}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f579939-6ba6-4397-b401-1e202fd36f50",
   "metadata": {},
   "source": [
    "## 取得存储在字典中的识别结果'text'字段"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "27c21299-0d04-430b-8fc1-0ba1ea117e81",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-----1-----\n",
      "数字仿真 \n",
      "类别\n",
      "数字字生 \n",
      "1．需求分析和场景定义：明确项目需求，如日照　1．数据采集：通过传感器、RFID 等技术实时采集物\n",
      "分析、结构性能、消防疏散等，并设定详细的仿\n",
      "理实体的运行数据，并进行数据清洗、整合和存储；\n",
      "真场;\n",
      "2 .建模与仿真：在虚拟空间中构建物理实体的数字\n",
      "2.  模型与优化：应用 BIM 软件构建三维模型,\n",
      "模型，并引入实际数据进行仿真，实现数据交互；\n",
      "并对模型进行必要的优化，如网格划分、边界条\n",
      "3．实时监测与预警：通过数字字生模型实时监测物\n",
      "件设置等；\n",
      "理实体的运行状态和性能变化, 一旦发现异常或潜在\n",
      "应用\n",
      "3．仿真参数设置：设定准确的仿真参数，如结构\n",
      "问题，及时发出预警信息；\n",
      "流程\n",
      "体系、栽荷取值、材料属性、环境条件等；\n",
      "4 .决策支持与优化：基于数字字生模型提供的数据\n",
      "4. 仿真分析：运行仿真程序，对结果进行可视化\n",
      "和分析结果，为决策者提供精准的决策支持，实现对\n",
      "处理分析，如应力分布、节能分析、流场等；\n",
      "物理\n",
      "5．成果反馈：并根据仿真结果及规范要求，优化\n",
      "设计方案。\n",
      "1. 软件: 根据项目需求, 选择专业仿真软件和分\n",
      "1.  软件：使用 loT 物联网平台、云计算、大数据分\n",
      "析软件，如 ANSYS、PKPM、Fluent、MATLAB\n",
      "析等先进技术来实现数字字生系统的构建和运行, 使\n",
      "软硬件 \n",
      "等进行性能分析和数据处理，使用 Revit，\n",
      "用专业 BIM 软件构建数字模型；\n",
      "Archicad 等 BIM 软件，进行三维建模；\n",
      "2. 硬件：传感器、数据采集设备、高性能服务器和\n",
      "工具\n",
      "2. 硬件：高性能计算机、工作站等，云计算或分\n",
      "存储设备等，用于数据采集、运算和存储，\n",
      "布式计算资源，\n",
      "1．模型精度：确保仿真模型能够准确反映项目实\n",
      "1．數据的实时性：确保数字字生模型能够实时反映\n",
      "际情况，包括结构、材料、环境等因素；\n",
      "物理\n",
      "2.  模型的更新与维护：随着物理实体的变化，及时\n",
      "2．参数设置：根据项目情况，合理设置参数，确\n",
      "技术\n",
      "保仿真结果准确；\n",
      "更新数字字生模型;\n",
      "要点\n",
      "3．结果验证：将仿真结果与规范标准及经验数据\n",
      "3.  数据安全与保护：在数据采集和传输过程中确保\n",
      "进行对比验证，确保结果可靠。\n",
      "数据的安全性和隐私性，防止数据泄露或被攻击。\n",
      "1．优化设计方案：通过仿真发现落在问题，提前\n",
      "1．提高运营效率：通过实时监测和控制运行状态，\n",
      "及时发现并解决问题，提高运营效率;\n",
      "进行优化；\n",
      "成果\n",
      "2.  降低维护成本：提前预测设备故障和维护需求，\n",
      "2.  提高产品质量：通过仿真验证产品性能，实现\n",
      "效果\n",
      "降本提质增效。\n",
      "实现植准维护和管理，降低维护成本。\n"
     ]
    }
   ],
   "source": [
    "for i in range(0, 1):\n",
    "    print('-----'+str(out_preds['Doc2'][i]['page'])+'-----')\n",
    "    \n",
    "    for polygon in out_preds['Doc2'][i]['text_lines']:\n",
    "        print(polygon['text'])\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d20fea6-6b19-4091-bb38-b857fb9da12a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
